use crate::{
    block_dev::BlockDevice,
    BlockID,
    BLOCK_SIZE,
};

use std::collections::HashMap;

pub struct BlockCacheDevice {
    block_device: Box<dyn BlockDevice>,
    caches: HashMap<BlockID, BlockCache>,
}

impl BlockDevice for BlockCacheDevice {
    fn read_block(&mut self, block_id: BlockID, buf: &mut [u8]) {
        let cache = self.get_block_cache(block_id);
        buf.copy_from_slice(cache.cache.as_slice());
    }

    fn write_block(&mut self, block_id: BlockID, buf: &[u8]) {
        let cache = self.get_block_cache(block_id);
        cache.modify = true;
        cache.cache.as_mut_slice().copy_from_slice(buf);
    }
}

impl Drop for BlockCacheDevice {
    fn drop(&mut self) {
        self.sync();
    }
}

impl BlockCacheDevice {
    pub fn new(block_device: Box<dyn BlockDevice>) -> Self {
        Self {
            block_device,
            caches: HashMap::new(),
        }
    }

    pub fn sync(&mut self) {
        for (block_id, cache) in self.caches.iter_mut() {
            if cache.modify {
                self.block_device.write_block(*block_id, cache.cache.as_slice());
                cache.modify = false;
            }
        }
    }

    pub fn get_ref<T>(&mut self, block_id: BlockID, offset: usize) -> &T 
        where T: Sized,
    {
        let type_size = core::mem::size_of::<T>();
        assert!(offset + type_size <= BLOCK_SIZE);

        let cache = self.get_block_cache(block_id);
        let ref_byte = &cache.cache[offset];
        unsafe {
            &*( ref_byte as *const _ as usize as *const T)
        }
    }

    pub fn get_mut<T>(&mut self, block_id: BlockID, offset: usize) -> &mut T 
        where T: Sized,
    {
        let type_size = core::mem::size_of::<T>();
        assert!(offset + type_size <= BLOCK_SIZE);

        let cache = self.get_block_cache(block_id);
        cache.modify = true;

        let ref_byte = &mut cache.cache[offset];
        unsafe {
            &mut *( ref_byte as *const _ as usize as *mut T)
        }
    }
    
}

impl BlockCacheDevice {
    fn get_block_cache(&mut self, block_id: BlockID) -> &mut BlockCache {
        self.caches.entry(block_id).or_insert_with(|| {
            let mut cache = Box::new([0u8; BLOCK_SIZE]);
            self.block_device.read_block(block_id, cache.as_mut_slice());
            BlockCache::new(cache)
        })
    }
}

struct BlockCache {
    cache: Box<[u8; BLOCK_SIZE]>,
    modify: bool,
}

impl BlockCache {
    fn new(cache: Box<[u8; BLOCK_SIZE]>) -> Self {
        Self {
            cache,
            modify: false,
        }
    }
}